// LOAD_WORK.cpp
// This module calculates concentrated loads work and acceleration field work
// Called by soluvdis in the original system

#include <vector>
#include <cmath>
#include <omp.h>

// Define real type based on precision needs
using Real = double;

// Global control parameters (would normally be in a configuration class)
namespace Control {
    extern Real DT2;       // Half time step
    extern Real XTWORK;    // Total work accumulator
    extern int NPROC;      // Number of processors
}

/**
 * @brief Calculates concentrated loads work and acceleration field work
 * 
 * This function computes:
 * 1) Work from concentrated loads
 * 2) Work from acceleration fields
 * 
 * @param nodePositions Node positions (3 coordinates per node)
 * @param nodeVelocities Node velocities (6 components per node)
 * @param loadNodes Load node information (5 parameters per load)
 * @param loadFactors Load scaling factors
 * @param forceValues Force values array
 * @param numLoads Number of concentrated loads
 * @param nodalMasses Nodal mass matrix (6 values per node)
 * @param accelFieldParams Acceleration field parameters (7 values per field)
 * @param accelFieldNodes Nodes in acceleration fields
 * @param numAccelFields Number of acceleration fields
 * @param frameVectors Reference frame vectors (16 values per frame)
 * @param globalVelocities Global velocities (3 components per node)
 */
void CalculateLoadWork(
    const std::vector<std::array<Real, 3>>& nodePositions,
    const std::vector<std::array<Real, 6>>& nodeVelocities,
    const std::vector<std::array<int, 5>>& loadNodes,
    const std::vector<Real>& loadFactors,
    const std::vector<Real>& forceValues,
    int numLoads,
    const std::vector<Real>& nodalMasses,
    const std::vector<std::array<Real, 7>>& accelFieldParams,
    const std::vector<int>& accelFieldNodes,
    int numAccelFields,
    const std::vector<std::array<Real, 16>>& frameVectors,
    const std::vector<std::array<Real, 3>>& globalVelocities)
{
    Real workIncrement = 0.0;

    // 1. Process concentrated loads
    for (int loadIdx = 0; loadIdx < numLoads; ++loadIdx) {
        const auto& load = loadNodes[loadIdx];
        Real force = forceValues[load[2]] * loadFactors[loadIdx] * Control::DT2;
        int direction = load[1];
        
        // Skip deactivated loads (negative direction)
        if (direction < 0) continue;
        
        int node = load[0];

        switch (direction) {
            case 1: // X direction
            case 2: // Y direction
            case 3: // Z direction
                workIncrement += force * nodeVelocities[node][direction - 1];
                break;
                
            case 4: { // Local frame direction
                int frame = load[3];
                workIncrement += force * (
                    frameVectors[frame][6] * nodeVelocities[node][0] +  // 7th element (0-based 6)
                    frameVectors[frame][7] * nodeVelocities[node][1] +  // 8th element
                    frameVectors[frame][8] * nodeVelocities[node][2]);   // 9th element
                break;
            }
                
            case 5: // Global X direction
            case 6: // Global Y direction
            case 7: // Global Z direction
                workIncrement += force * globalVelocities[node][direction - 5];
                break;
                
            case 8: { // Global frame direction
                int frame = load[3];
                workIncrement += force * (
                    frameVectors[frame][6] * globalVelocities[node][0] +
                    frameVectors[frame][7] * globalVelocities[node][1] +
                    frameVectors[frame][8] * globalVelocities[node][2]);
                break;
            }
        }
    }

    // 2. Process acceleration fields
    int nodeOffset = 0;
    std::vector<Real> procWorkIncrements(Control::NPROC, 0.0);

    for (int fieldIdx = 0; fieldIdx < numAccelFields; ++fieldIdx) {
        const auto& field = accelFieldParams[fieldIdx];
        int numBoundaryNodes = static_cast<int>(field[0]);
        
        if (numBoundaryNodes == 0) continue;
        
        int xSensor = static_cast<int>(field[1]);
        int ySensor = static_cast<int>(field[2]);
        int zSensor = static_cast<int>(field[3]);
        Real xAccel = field[4];
        Real yAccel = field[5];
        Real zAccel = field[6];

        int nodesPerProc = (numBoundaryNodes + Control::NPROC - 1) / Control::NPROC;
        std::fill(procWorkIncrements.begin(), procWorkIncrements.end(), 0.0);

        // Process X acceleration component
        if (xSensor != 0) {
            Real acceleration = xAccel * forceValues[xSensor] * Control::DT2;
            
            #pragma omp parallel for
            for (int proc = 0; proc < Control::NPROC; ++proc) {
                int startNode = proc * nodesPerProc;
                int endNode = std::min(numBoundaryNodes, (proc + 1) * nodesPerProc);
                
                for (int i = startNode; i < endNode; ++i) {
                    int node = accelFieldNodes[nodeOffset + i];
                    Real force = nodalMasses[6 * node] * acceleration; // X mass component
                    procWorkIncrements[proc] += force * nodeVelocities[node][0];
                }
            }
        }

        // Process Y acceleration component
        if (ySensor != 0) {
            Real acceleration = yAccel * forceValues[ySensor] * Control::DT2;
            
            #pragma omp parallel for
            for (int proc = 0; proc < Control::NPROC; ++proc) {
                int startNode = proc * nodesPerProc;
                int endNode = std::min(numBoundaryNodes, (proc + 1) * nodesPerProc);
                
                for (int i = startNode; i < endNode; ++i) {
                    int node = accelFieldNodes[nodeOffset + i];
                    Real force = nodalMasses[6 * node + 1] * acceleration; // Y mass component
                    procWorkIncrements[proc] += force * nodeVelocities[node][1];
                }
            }
        }

        // Process Z acceleration component
        if (zSensor != 0) {
            Real acceleration = zAccel * forceValues[zSensor] * Control::DT2;
            
            #pragma omp parallel for
            for (int proc = 0; proc < Control::NPROC; ++proc) {
                int startNode = proc * nodesPerProc;
                int endNode = std::min(numBoundaryNodes, (proc + 1) * nodesPerProc);
                
                for (int i = startNode; i < endNode; ++i) {
                    int node = accelFieldNodes[nodeOffset + i];
                    Real force = nodalMasses[6 * node + 2] * acceleration; // Z mass component
                    procWorkIncrements[proc] += force * nodeVelocities[node][2];
                }
            }
        }

        // Accumulate work from all processors
        for (Real procWork : procWorkIncrements) {
            workIncrement += procWork;
        }

        nodeOffset += numBoundaryNodes;
    }

    // Update total work
    Control::XTWORK += workIncrement;
}